import subprocess

from netmiko import ConnectHandler

from interface import Interface
from tools import format_cli_output, firewall_policy_name
from tools import normalize_route_info, normalize_interface_ip_index, normalize_interface_index_name, \
    normalize_interface_ip_netmask, normalize_interface_info, normalize_route_info_interface_index, cidr_to_netmask


class Firewall:

    def __init__(self, device_type, device_name, manage_ip, auth_type, snmp_port, snmp_password, ssh_port, ssh_username,
                 ssh_password):
        """
        初始化Router类的实例。

        :param device_type: 设备类型
        :param device_name: 设备名称
        :param manage_ip: 管理IP地址
        :param auth_type: 认证类型（如SNMP，SSH等）
        :param snmp_port: SNMP端口号
        :param snmp_password: SNMP密码
        :param ssh_port: SSH端口号
        :param ssh_username: SSH用户名
        :param ssh_password: SSH密码
        """
        self.device_type = device_type  # 设备类型
        self.device_name = device_name  # 设备名称
        self.manage_ip = manage_ip  # 管理IP地址
        self.auth_type = auth_type  # 认证类型
        self.snmp_port = snmp_port  # SNMP端口号
        self.snmp_password = snmp_password  # SNMP密码
        self.ssh_port = ssh_port  # SSH端口号
        self.ssh_username = ssh_username  # SSH用户名
        self.ssh_password = ssh_password  # SSH密码

        self.secure_zone = self.get_secure_zone()
        self.route_info_proto = self.get_ip_route_table_proto()
        self.route_info_interface = self.get_ip_route_table_interface()
        self.route_info = self.combine_route_table()

        self.interface_ip_index_info = self.get_interface_ip_index_info()
        self.interface_index_name_info = self.get_interface_index_name_info()
        self.interface_ip_netmask_info = self.get_interface_ip_netmask_info()
        self.interface_info = self.combine_interface_info()

        self.secure_policy_info = self.get_secure_policy()

    # 合并路由信息，获得最终路由表
    def combine_route_table(self):
        route_info_proto = self.route_info_proto
        route_info_interface = self.route_info_interface
        route_info_combined = []
        for route_info_proto_element in route_info_proto:
            for route_info_interface_element in route_info_interface:
                if route_info_proto_element["destination"] == route_info_interface_element["destination"]:
                    route_info_proto_element["interface"] = route_info_interface_element["interface"]
                    route_info_combined.append(route_info_proto_element)
        print("获得设备{}路由表:".format(self.device_name))
        format_cli_output(rows=[[[route_info_combined_element["destination"]],
                                 [route_info_combined_element["net_mask"]],
                                 [route_info_combined_element["nexthop"]],
                                 [route_info_combined_element["interface"]],
                                 [route_info_combined_element["proto_type"]]] for route_info_combined_element in
                                route_info_combined],
                          headers=["destination", "net_mask", "nexthop", "interface", "proto_type"])

        return route_info_combined

    # snmp获取路由——proto信息
    def get_ip_route_table_proto(self):
        route_table = []
        # print("获取到设备 {} 的路由信息".format(self.device_name))
        snmp_command = './snmptool/snmpwalk.exe -v 2c -c {} {} {}' \
            .format(self.snmp_password, self.manage_ip, "1.3.6.1.2.1.4.24.4.1.6")
        completed_process = subprocess.Popen(snmp_command, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout, stderr = completed_process.communicate()
        if len(stdout) > 0:
            snmp_message = stdout.decode()
            route_table = normalize_route_info(snmp_message)
            # print("设备 {} 的路由信息为：\n{}".format(self.device_name, route_table))
        elif len(stderr) > 0:
            print("设备 {} 的路由信息获取错误，错误信息为： {}".format(self.device_name, stderr.decode()))
            raise Exception("设备 {} 的接口信息获取错误，错误信息为： {}".format(self.device_name, stderr.decode()))
        return route_table

    # snmp获取路由——proto信息
    def get_ip_route_table_interface(self):
        route_table = []
        # print("获取到设备 {} 的路由信息".format(self.device_name))
        snmp_command = './snmptool/snmpwalk.exe -v 2c -c {} {} {}' \
            .format(self.snmp_password, self.manage_ip, "1.3.6.1.2.1.4.24.4.1.5")
        completed_process = subprocess.Popen(snmp_command, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout, stderr = completed_process.communicate()
        if len(stdout) > 0:
            snmp_message = stdout.decode()
            route_table = normalize_route_info_interface_index(snmp_message)
            # print("设备 {} 的路由信息为：\n{}".format(self.device_name, route_table))
        elif len(stderr) > 0:
            print("设备 {} 的路由信息获取错误，错误信息为： {}".format(self.device_name, stderr.decode()))
            raise Exception("设备 {} 的接口信息获取错误，错误信息为： {}".format(self.device_name, stderr.decode()))
        return route_table

    # snmp获取接口index信息，格式为ip = index
    #  示例： .1.3.6.1.2.1.4.20.1.2.10.0.10.2 = INTEGER: 4
    def get_interface_ip_index_info(self):
        # print("获取到设备 {} 的接口信息".format(self.device_name))
        interface_ip_index_list = []
        snmp_command = './snmptool/snmpwalk.exe -v 2c -c {} {} {}' \
            .format(self.snmp_password, self.manage_ip, "1.3.6.1.2.1.4.20.1.2")
        completed_process = subprocess.Popen(snmp_command, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout, stderr = completed_process.communicate()
        if len(stdout) > 0:
            snmp_message = stdout.decode()
            interface_ip_index_list = normalize_interface_ip_index(snmp_message)
            # print("设备 {} 的接口信息为： {}".format(self.device_name, interface_ip_index_list))
        elif len(stderr) > 0:
            print("设备 {} 的接口信息获取错误，错误信息为： {}".format(self.device_name, stderr.decode()))
            raise Exception("设备 {} 的接口信息获取错误，错误信息为： {}".format(self.device_name, stderr.decode()))
        return interface_ip_index_list

    # snmp获取接口index信息，格式为 index = name
    # 示例： .1.3.6.1.2.1.31.1.1.1.18.4 = STRING: "HUAWEI, AR Series, GigabitEthernet0/0/1 Interface"
    def get_interface_index_name_info(self):
        # print("获取到设备 {} 的接口信息".format(self.device_name))
        interface_index_name_list = []
        snmp_command = './snmptool/snmpwalk.exe -v 2c -c {} {} {}' \
            .format(self.snmp_password, self.manage_ip, "1.3.6.1.2.1.31.1.1.1.18")
        completed_process = subprocess.Popen(snmp_command, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout, stderr = completed_process.communicate()
        if len(stdout) > 0:
            snmp_message = stdout.decode()
            interface_index_name_list = normalize_interface_index_name(snmp_message)
            # print("设备 {} 的接口信息为： {}".format(self.device_name, interface_index_name_list))
        elif len(stderr) > 0:
            print("设备 {} 的接口信息获取错误，错误信息为： {}".format(self.device_name, stderr.decode()))
            raise Exception("设备 {} 的接口信息获取错误，错误信息为： {}".format(self.device_name, stderr.decode()))
        return interface_index_name_list

    # snmp获取接口index信息，格式为 ip = netmask
    # 示例： .1.3.6.1.2.1.4.20.1.3.10.0.10.2 = IpAddress: 255.255.255.252
    def get_interface_ip_netmask_info(self):
        # print("获取到设备 {} 的接口信息".format(self.device_name))
        interface_ip_netmask_list = []
        snmp_command = './snmptool/snmpwalk.exe -v 2c -c {} {} {}' \
            .format(self.snmp_password, self.manage_ip, "1.3.6.1.2.1.4.20.1.3")
        completed_process = subprocess.Popen(snmp_command, shell=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        stdout, stderr = completed_process.communicate()
        if len(stdout) > 0:
            snmp_message = stdout.decode()
            interface_ip_netmask_list = normalize_interface_ip_netmask(snmp_message)
            # print("设备 {} 的接口信息为： {}".format(self.device_name, interface_ip_netmask_list))
        elif len(stderr) > 0:
            print("设备 {} 的接口信息获取错误，错误信息为： {}".format(self.device_name, stderr.decode()))
            raise Exception("设备 {} 的接口信息获取错误，错误信息为： {}".format(self.device_name, stderr.decode()))
        return interface_ip_netmask_list

    # todo 完善
    # 获取防火墙安全策略信息
    def get_secure_policy(self):
        # ssh连接获取安全策略信息
        ssh_secure_policy_commands = ['security-policy', 'dis this']
        security_policy_text = self.execute_secure_policy_command(ssh_secure_policy_commands, terminator=r"return")
        secure_policy = []
        current_rule = None
        lines = [line for line in security_policy_text.strip().split("\n") if line.strip()]
        for line in lines:
            if "rule name" in line:
                rule_name = line.split()[-1].strip()
                current_rule = {
                    "rule name": rule_name,
                    "source-zone": [],
                    "destination-zone": [],
                    "source-address": [],
                    "destination-address": [],
                    "action": ""
                }
                secure_policy.append(current_rule)
            elif "source-zone" in line:
                source_zone = line.split()[-1].strip()
                current_rule["source-zone"].append(source_zone)
            elif "destination-zone" in line:
                destination_zone = line.split()[-1].strip()
                current_rule["destination-zone"].append(destination_zone)
            elif "source-address" in line:
                parts = line.split()
                if len(parts) > 3:
                    ip, mask = parts[1], parts[3]
                    current_rule["source-address"].append({"ip": ip, "mask": mask})
            elif "destination-address" in line:
                parts = line.split()
                if len(parts) > 3:
                    ip, mask = parts[1], parts[3]
                    current_rule["destination-address"].append({"ip": ip.strip(), "mask": mask.strip()})
            elif "action" in line:
                action = line.split()[-1].strip()
                current_rule["action"] = action
        # 示例，安全策略原格式：
        '''
        [USG6000V1]security-policy 
#
        USG6000V1-policy-security]dis this
        024-03-07 05:04:28.470 
        
        security-policy
        rule name dmz2untrust
         source-zone dmz
         destination-zone untrust
         source-address 10.0.1.0 mask 255.255.255.0
         source-address 10.0.2.0 mask 255.255.255.0
         destination-address 172.16.6.0 mask 255.255.255.0
         action permit
        rule name manage2all
         source-zone trust
         source-address 192.168.12.0 mask 255.255.255.0
         action permit
#
        '''
        # 安全策略格式化后：
        '''
        secure_policy =[
        {
        "rule name":"dmz2untrust",
        "source-zone":["dmz"],
        "source-address":[{"ip":"10.0.1.0","mask:"255.255.255.0"},{"ip":"10.0.2.0","mask:"255.255.255.0"}],
        "destination-zone":["untrust"],
        "destination-address":[{"ip":"172.16.6.0","mask":"255.255.255.0"}]
        "action":"permit"
        },
        {
        "rule name":"manage2all",
        "source-zone":["trust"],
        "source-address":[{"ip":"192.168.12.0","mask:"255.255.255.0"}],
        "destination-zone":[],
        "destination-address":[]
        "action":"permit"
        }
        ]
        '''
        print("防火墙{}安全策略信息：".format(self.device_name))
        for secure_policy_element in secure_policy:
            print(
                "策略名：{}\n  source-zone:{}\n  source-address:{}\n  destination-zone:{}\n  destination-address:{}\n  "
                "action:{} "
                    .format(secure_policy_element["rule name"] if "rule name" in secure_policy_element else "null",
                            secure_policy_element["source-zone"] if "source-zone" in secure_policy_element else "null",
                            str(secure_policy_element[
                                    "source-address"]) if "source-address" in secure_policy_element else "null",
                            secure_policy_element[
                                "destination-zone"] if "destination-zone" in secure_policy_element else "null",
                            str(secure_policy_element[
                                    "destination-address"]) if "destination-address" in secure_policy_element else "null",
                            secure_policy_element["action"] if "action" in secure_policy_element else "null"))
        return secure_policy

    # 获取防火墙安全区域信息
    def get_secure_zone(self):
        # ssh连接获取安全区域信息
        ssh_secure_zone_commands = ['security-policy', 'dis zone']
        secure_zone_text = self.execute_secure_policy_command(ssh_secure_zone_commands, terminator=r"return")
        zone_info = []
        blocks = secure_zone_text.split("#")
        for block in blocks:
            if ("priority" in block):
                lines = [line for line in block.strip().split("\n") if line.strip()]
                index = None
                try:
                    for idx in range(len(lines)):
                        if ("priority" in lines[idx]):
                            index = idx
                except ValueError:
                    print(f"无元素")
                if index is not None:
                    number_string = lines[index + 1][lines[index + 1].index("(") + 1:lines[index + 1].index(")")]
                    number = int(number_string)
                    interface = []
                    current_rule = {
                        'zone_name': lines[index - 1].split()[-1].strip(),
                        lines[index].split()[0].strip(): lines[index].split()[-1].strip(),
                        lines[index + 1].split()[0].strip(): interface
                    }
                    for x in range(number):
                        interface.append(lines[index + 2 + x].strip())
                    zone_info.append(current_rule)
        print("防火墙{}安全区域信息：".format(self.device_name))
        for zone_info_element in zone_info:
            print(zone_info_element)
        return zone_info

        # 示例，安全区域原格式：
        '''
        [USG6000V1]dis zone
        2024-03-07 05:12:11.850 
        local
         priority is 100
         interface of the zone is (0):
        #
        trust
         priority is 85
         interface of the zone is (4):
            GigabitEthernet0/0/0
            GigabitEthernet1/0/0
            GigabitEthernet1/0/1
            GigabitEthernet1/0/2
        #
        untrust
         priority is 5
         interface of the zone is (1):
            GigabitEthernet1/0/4
        #
        dmz
         priority is 50
         interface of the zone is (1):
            GigabitEthernet1/0/3
        #
        '''
        pass

    def create_secure_policy(self, secure_route_info_detail, origin_ip, destination_ip):
        interface_in = secure_route_info_detail["interface_in"]
        interface_out = secure_route_info_detail["interface_out"]
        print("设备 {} 正在分析中，新建防火墙策略：".format(self.device_name))
        source_zone = interface_in.zone
        source_address = origin_ip.split("/")[0]
        source_mask = origin_ip.split("/")[1]
        destination_zone = interface_out.zone
        destination_address = destination_ip.split("/")[0]
        destination_mask = destination_ip.split("/")[1]
        # 防火墙名随机
        rule_name = firewall_policy_name(source_zone, destination_zone)
        action = "permit"
        secure_policy_command = ["security-policy",
                                 "rule name {}".format(rule_name),
                                 "source-zone {}".format(source_zone),
                                 "source-address {} mask {}".format(source_address, source_mask),
                                 "destination-zone {}".format(destination_zone),
                                 "destination-address {} mask {}".format(destination_address, destination_mask),
                                 "action {}".format(action),
                                 "dis this"]
        print(secure_policy_command)
        return secure_policy_command

    def delete_secure_policy(self):
        return

    # 修改安全策略，目前仅支持扩容origin-address和destination-address
    def edit_secure_policy(self, old_policy, edit_key, origin_ip, destination_ip):
        print("设备 {} 正在分析中，修改防火墙策略：".format(self.device_name))
        rule_name = old_policy["rule name"]
        source_zone = old_policy["source-zone"]
        source_address = old_policy["source-address"]
        destination_zone = old_policy["destination-zone"]
        source_address = old_policy["destination-address"]
        action = old_policy["action"]
        if edit_key == "source-address":
            secure_policy_command = ["security-policy",
                                     "rule name {}".format(rule_name),
                                     "source-address {} mask {}".format(origin_ip.split("/")[0],
                                                                        origin_ip.split("/")[1]),
                                     "dis this"]
        elif edit_key == "destination-address":
            secure_policy_command = ["security-policy",
                                     "rule name {}".format(rule_name),
                                     "destination-address {} mask {}".format(destination_ip.split("/")[0],
                                                                             destination_ip.split("/")[1]),
                                     "action {}".format(action),
                                     "dis this"]
        else:
            secure_policy_command = ["security-policy", "dis this"]
        # print(secure_policy_command)
        return secure_policy_command

    # 合并所有接口信息
    def combine_interface_info(self):
        interface_ip_index_info = self.interface_ip_index_info
        interface_index_name_info = self.interface_index_name_info
        interface_ip_netmask_info = self.interface_ip_netmask_info
        interface_info_temp = []
        interface_info = []
        for interface_ip_index_info_element in interface_ip_index_info:
            for interface_index_name_info_element in interface_index_name_info:
                if interface_ip_index_info_element["index"] == interface_index_name_info_element["index"]:
                    interface_ip_index_info_element["name"] = interface_index_name_info_element["name"]
                    interface_info_temp.append(interface_ip_index_info_element)
        for interface_info_element in interface_info_temp:
            for interface_ip_netmask_info_element in interface_ip_netmask_info:
                if interface_info_element["ip"] == interface_ip_netmask_info_element["ip"]:
                    interface_info.append(Interface(interface_info_element["index"], interface_info_element["name"],
                                                    interface_info_element["ip"],
                                                    interface_ip_netmask_info_element["net_mask"]))
        # 补全接口中zone的信息
        for interface in interface_info:
            for zone in self.secure_zone:
                zone_name = zone["zone_name"]
                zone_contain_interface_list = zone["interface"]
                for zone_contain_interface in zone_contain_interface_list:
                    if zone_contain_interface in interface.name:
                        interface.set_zone(zone_name)
        print("设备{}接口（使用中）：".format(self.device_name))
        for interface in interface_info:
            interface.display_info()
        return interface_info

    # 执行网络策略相关命令
    def execute_secure_policy_command(self, commands, terminator=r"#"):
        ssh_device = {
            'device_type': 'huawei_vrp',
            'host': self.manage_ip,
            'username': self.ssh_username,
            'password': self.ssh_password,
            'port': self.ssh_port,  # optional, defaults to 22
            'secret': '',  # optional, defaults to ''
        }
        with ConnectHandler(**ssh_device) as connect:
            # commands = ['security-policy', 'rule name dmz2untrust',
            #             'dis this']
            output = connect.send_config_set(commands, exit_config_mode=True, terminator=terminator)
        print("防火墙执行了命令：{}".format(str(commands)))
        return output

    def print_device_details(self):
        print("读取到设备{}的详细信息如下：".format(self.device_name))
        print("读取到防火墙{}接口（使用中）：".format(self.device_name))
        for interface in self.interface_info:
            interface.display_info()

        print("读取到防火墙{}路由表:".format(self.device_name))
        format_cli_output(rows=[[[route_info_combined_element["destination"]],
                                 [route_info_combined_element["net_mask"]],
                                 [route_info_combined_element["nexthop"]],
                                 [route_info_combined_element["interface"]],
                                 [route_info_combined_element["proto_type"]]] for route_info_combined_element in
                                self.route_info],
                          headers=["destination", "net_mask", "nexthop", "interface", "proto_type"])
        print("读取到防火墙{}安全区域信息：".format(self.device_name))
        for secure_policy_zone in self.secure_zone:
            print(secure_policy_zone)
        print("读取到防火墙{}安全策略信息：".format(self.device_name))
        for secure_policy_element in self.secure_policy_info:
            print(
                "策略名：{}\n  source-zone:{}\n  source-address:{}\n  destination-zone:{}\n  destination-address:{}\n  "
                "action:{} "
                    .format(secure_policy_element["rule name"] if "rule name" in secure_policy_element else "null",
                            secure_policy_element["source-zone"] if "source-zone" in secure_policy_element else "null",
                            str(secure_policy_element[
                                    "source-address"]) if "source-address" in secure_policy_element else "null",
                            secure_policy_element[
                                "destination-zone"] if "destination-zone" in secure_policy_element else "null",
                            str(secure_policy_element[
                                    "destination-address"]) if "destination-address" in secure_policy_element else "null",
                            secure_policy_element["action"] if "action" in secure_policy_element else "null"))


if __name__ == '__main__':
    ustr = {"commands": ["security-policy",
                         "rule name example_rule",
                         "source-zone internal",
                         "source-address 192.168.1.0 mask 255.255.255.0",
                         "destination-zone external",
                         "destination-address 10.0.0.0 mask 255.0.0.0",
                         "action permit",
                         "dis this"], "edit_key": "source-address"}
    cstr = {"commands": ["security-policy",
                         "rule name example_rule",
                         "source-zone internal",
                         "source-address 192.168.1.0 mask 255.255.255.0",
                         "destination-zone external",
                         "destination-address 10.0.0.0 mask 255.0.0.0",
                         "action permit",
                         "dis this"]}

    pass
